package com.carp.shard.plugin.processor;


import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
import com.alibaba.druid.stat.TableStat;
import com.alibaba.fastjson.JSON;
import com.carp.shard.plugin.ShardPlugin;
import com.carp.shard.plugin.constant.ReadType;
import com.carp.shard.plugin.constant.WriteType;
import com.carp.shard.plugin.exception.MyBatisShardException;
import com.carp.shard.plugin.strategy.AbstractShardStrategy;
import com.carp.shard.plugin.strategy.ShardStrategyContext;
import com.carp.shard.plugin.utils.CommonUtils;
import com.carp.shard.plugin.visitor.MySqlReplaceTableNameVisitor;
import lombok.Getter;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.parameter.ParameterHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.scripting.defaults.DefaultParameterHandler;
import org.springframework.util.Assert;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.util.List;
import java.util.Map;

/**
 * 分片处理器
 *
 * @author lile
 * @since 2024/5/11
 */
@Slf4j
@Getter
public class ShardProcessor {
    private final MetaObject metaObject;
    private final BoundSql boundSql;
    private final MappedStatement mappedStatement;
    private final String originalSql;
    private final String shardSql;
    private boolean replacedTableName = false;
    private ReadType readType = ReadType.OLD;
    private WriteType writeType = WriteType.OLD;

    public ShardProcessor(@NonNull MetaObject metaObject) {
        this.metaObject = metaObject;
        this.boundSql = (BoundSql) metaObject.getValue("delegate.boundSql");
        this.mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");
        this.originalSql = boundSql.getSql();
        this.shardSql = calculateShardSql();
    }

    /**
     * 处理参数
     */
    public void processParams() {
        SQLStatement stmt = getSqlStatement();
        // 统计SQL中使用的表、字段、过滤条件、排序表达式、分组表达式等
        SchemaStatVisitor schemaStatVisitor = SQLUtils.createSchemaStatVisitor(ShardPlugin.DB_TYPE);
        stmt.accept(schemaStatVisitor);
        // 表名
        Map<TableStat.Name, TableStat> tables = schemaStatVisitor.getTables();

        for (TableStat.Name name : tables.keySet()) {
            AbstractShardStrategy shardStrategy = ShardStrategyContext.getStrategyByTableName(name.getName());
            shardStrategy.processParams(metaObject, boundSql, schemaStatVisitor);
        }
    }

    /**
     * 路由到指定分片
     */
    public void route() {
        if (readType == ReadType.OLD) {
            return;
        }
        setShardSql();
    }

    /**
     * 处理写请求
     *
     * @param connection 数据库连接
     */
    public void processWrite(@NonNull Connection connection) {
        if (!replacedTableName) {
            return;
        }

        switch (writeType) {
            case OLD:
                return;
            case NEW:
                setShardSql();
                return;
            case BOTH:
                writeShard(connection);
                return;
            default:
                throw new MyBatisShardException("未知ReadType：" + readType);
        }

    }

    /**
     * 写分片
     *
     * @param connection 数据库连接
     */
    private void writeShard(Connection connection) {
        Object parameterObject = boundSql.getParameterObject();
        log.info("[shard-plugin] double write sql: {} \n  parameterObject({}): {}", CommonUtils.removeBreakingWhitespace(shardSql), parameterObject.getClass().getSimpleName(), JSON.toJSONString(parameterObject));
        try (PreparedStatement statement = connection.prepareStatement(shardSql)) {
            ParameterHandler parameterHandler = new DefaultParameterHandler(
                    mappedStatement, parameterObject, boundSql);
            parameterHandler.setParameters(statement);
            statement.executeUpdate();
        } catch (Exception e) {
            throw new MyBatisShardException(String.format("Error: Method ShardPlugin.write execution error of sql : \n %s \n", CommonUtils.removeBreakingWhitespace(shardSql)), e);
        }
    }

    /**
     * 将SQL设置为分片后的SQL
     */
    private void setShardSql() {
        metaObject.setValue("delegate.boundSql.sql", shardSql);
    }

    /**
     * 获取 AST
     *
     * @return AST
     */
    private SQLStatement getSqlStatement() {
        List<SQLStatement> stmtList = SQLUtils.parseStatements(originalSql, ShardPlugin.DB_TYPE);
        Assert.notEmpty(stmtList, "stmtList is empty, sql: " + originalSql);
        return stmtList.get(0);
    }

    /**
     * 计算分片后的SQL
     *
     * @return 分片后的SQL
     */
    private String calculateShardSql() {
        SQLStatement stmt = getSqlStatement();
        MySqlReplaceTableNameVisitor mySqlReplaceTableNameVisitor = new MySqlReplaceTableNameVisitor(boundSql);
        stmt.accept(mySqlReplaceTableNameVisitor);
        setFlag(mySqlReplaceTableNameVisitor);
        return SQLUtils.toMySqlString(stmt);
    }

    /**
     * 设置标志位
     *
     * @param mySqlReplaceTableNameVisitor 用于替换表名的 Visitor
     */
    private void setFlag(MySqlReplaceTableNameVisitor mySqlReplaceTableNameVisitor) {
        this.replacedTableName = mySqlReplaceTableNameVisitor.isAnyReplaced();
        this.readType = mySqlReplaceTableNameVisitor.getFinalReadType();
        this.writeType = mySqlReplaceTableNameVisitor.getFinalWriteType();
    }

}
